import akshare as ak
import pandas as pd
import time
import os
import random
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

import pymysql
from fake_useragent import UserAgent




def get_stock_data(symbol, stock_name,start_date, end_date, retry=1000):
    """获取股票历史数据（增加虚拟请求头防封禁）"""
    for attempt in range(retry):
        try:
            print(f"🔄 正在获取 {symbol} 数据 (尝试 {attempt + 1}/{retry})...")



            # 获取行情数据（使用带虚拟头的请求）
            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period="daily",
                start_date=start_date,
                end_date=end_date,
                adjust="qfq"
            )

            # 检查数据是否为空
            if df.empty:
                raise ValueError("获取的数据为空")

            # 列重命名
            df = df.rename(columns={
                '日期': '日期',
                '开盘': '开盘价',
                '收盘': '收盘价',
                '最高': '最高价',
                '最低': '最低价',
                '成交量': '成交量(手)',
                '成交额': '成交额(元)',
                '涨跌幅': '涨跌幅(%)',
                '换手率': '换手率(%)',
                '振幅': '振幅(%)'
            })

            # 数据处理
            df['涨跌额'] = (df['收盘价'] - df['开盘价'])
            # df['涨跌幅(%)'] = ((df['涨跌额'] / df['收盘价'])*100).round(2)
            df['股票代码'] = symbol
            df['股票名称'] = stock_name
            df['成交额(亿元)'] = (df['成交额(元)'] / 1e8).round(2)
            df['日期'] = pd.to_datetime(df['日期']).dt.strftime('%Y-%m-%d')


            return df[[
                '日期', '股票代码', '股票名称', '开盘价', '收盘价', '最高价', '最低价',
                '成交量(手)', '成交额(亿元)', '涨跌幅(%)', '涨跌额', '换手率(%)', '振幅(%)'
            ]]



        except Exception as e:
            print(f"⚠️ 第{attempt + 1}次尝试失败: {str(e)}")


    raise ConnectionError("❌ 数据获取失败，请检查参数和网络连接")


def save_to_excel(data, symbol, start_date, end_date):
    """保存数据到Excel（专业格式化）"""
    os.makedirs("../gp_xlsx", exist_ok=True)
    file_name = os.path.join("../gp_xlsx", f"{symbol}_{start_date}_{end_date}.xlsx")

    with pd.ExcelWriter(file_name, engine='openpyxl') as writer:
        data.to_excel(writer, sheet_name='股票数据', index=False)
        workbook = writer.book
        worksheet = writer.sheets['股票数据']

        # 设置专业格式
        for row in range(2, len(data) + 2):
            if '成交额(亿元)' in data.columns:
                worksheet.cell(row=row, column=data.columns.get_loc('成交额(亿元)') + 1).number_format = '0.00""'
            if '涨跌幅(%)' in data.columns:
                worksheet.cell(row=row, column=data.columns.get_loc('涨跌幅(%)') + 1).number_format = '0.00'
            if '涨跌额' in data.columns:
                worksheet.cell(row=row, column=data.columns.get_loc('涨跌额') + 1).number_format = '0.00;-0.00;0'

    print(f"📁 文件已保存: {os.path.abspath(file_name)}")
    return os.path.abspath(file_name)


def is_file_exists(symbol, start_date, end_date):
    """检查股票文件是否已存在"""
    file_name = os.path.join("../gp_xlsx", f"{symbol}_{start_date}_{end_date}.xlsx")
    return os.path.exists(file_name)


def process_single_stock_without_check(symbol, stock_name,start_date, end_date):
    """处理单只股票的数据获取和保存（增加文件存在检查）"""
    try:
        # 首先检查文件是否已存在
        if is_file_exists(symbol, start_date, end_date):
            print(f"⏩ 股票 {symbol} 文件已存在，跳过处理")
            return symbol, "skipped", f"{symbol}_{start_date}_{end_date}.xlsx"

        print(f"⏳ 开始处理股票 {symbol} ...")
        stock_data = get_stock_data(symbol, stock_name,start_date, end_date)
        file_path = save_to_excel(stock_data, symbol, start_date, end_date)
        # print(f"✅ {symbol} 处理完成: {file_path}")
        return symbol, "success", file_path
    except Exception as e:
        print(f"❌ 股票 {symbol} 处理失败: {str(e)}")
        return symbol, "failed", str(e)

def get_stock_codes():
    default_db_config = {
        'host': '175.178.5.14',
        'port': 33066,
        'user': 'root',
        'password': '07M@Y0sqL1',
        'database': 'stock',  # 假设数据库名为stock，如果不是，请修改
        'charset': 'utf8mb4',
        'connect_timeout': 10  # 添加连接超时设置
    }

    # 如果提供了自定义配置，则使用自定义配置
    # if db_config:
    #     for key, value in db_config.items():
    #         if value:  # 只更新非空值
    #             default_db_config[key] = value

    db_config = default_db_config

    # 构建SQL查询中的IN子句

    # 构建完整的SQL查询
    sql = "select SYMBOL,NAME from  stock_personal_data_info"

    print(f"执行SQL查询: {sql}")

    # 连接数据库并执行查询
    sectors = {}
    try:
        print("尝试连接数据库...")
        conn = pymysql.connect(**db_config)
        print("数据库连接成功")

        with conn.cursor() as cursor:
            print("执行SQL查询...")
            cursor.execute(sql)
            results = cursor.fetchall()
            print(f"查询结果: {len(results)} 条记录")



            for symbol,name in results:
                sectors[symbol] = name
    except pymysql.err.OperationalError as e:
        print(f"数据库连接失败: {e}")
        print("将继续处理数据，但不包含板块信息")
    except Exception as e:
        print(f"数据库查询失败: {e}")
        print("将继续处理数据，但不包含板块信息")
    finally:
        if 'conn' in locals() and conn:
            conn.close()
            print("数据库连接已关闭")
    return sectors


def main():
    """主程序入口（多线程批量处理模式）"""
    print("=" * 50)
    print("📊 股票历史数据批量获取工具（多线程版）")
    print("=" * 50)
    print("💡 提示：从gp_data/stock_codes.xlsx读取股票代码列表")
    print(f"🔒 安全机制：启用虚拟请求头防封禁（随机User-Agent+Referer）")

    # 读取Excel文件中的股票代码
    # excel_path = "../gp_data/gp_code_20250731.xls"
    # try:
    #     df_codes = pd.read_excel(excel_path, sheet_name=0)
    #     code_col = "成份券代码Constituent Code"
    #     name_col ="成份券名称Constituent Name"
    #     stock_codes = df_codes[code_col].astype(str).str.zfill(6).tolist()
    #     stock_df = pd.DataFrame({
    #         '股票代码': df_codes[code_col].astype(str).str.zfill(6),
    #         '股票名称': df_codes[name_col].astype(str)  # 假设name_col是股票名称列
    #     })
    #     print(f"🔍 从Excel读取到 {len(stock_codes)} 个股票代码")
    # except Exception as e:
    #     print(f"❌ 读取股票代码文件失败: {str(e)}")
    #     print("💡 请确保文件路径正确，且包含股票代码列")
    #     return

    #从数据库读取股票代码
    stock_dict = get_stock_codes()
    stock_codes = list(stock_dict.keys())
    stock_df = pd.DataFrame({
        '股票代码': list(stock_dict.keys()),
        '股票名称': list(stock_dict.values())
    })


    # 用户输入日期范围
    start_date = '19900101'
    end_date = input(f"结束日期（默认今天）: ").strip() or datetime.now().strftime("%Y%m%d")

    # 日期格式校验
    if len(start_date) != 8 or not start_date.isdigit():
        print("❌ 开始日期格式应为YYYYMMDD")
        return
    if len(end_date) != 8 or not end_date.isdigit():
        end_date = datetime.now().strftime("%Y%m%d")

    start = time.perf_counter()
    # 记录不同状态的股票
    success_list = []
    skipped_list = []
    fail_list = []

    # === 关键改造：提前批量检查文件存在性 ===
    # 一次性计算需要处理的股票代码（文件不存在的）
    to_process = []
    for symbol in stock_codes:
        if is_file_exists(symbol, start_date, end_date):
            print(f"⏩ 股票 {symbol} 文件已存在，跳过处理")
            skipped_list.append(symbol)
        else:
            to_process.append(symbol)
            print(f"⏳ 股票 {symbol} 需要处理")

    # 获取需要移除的股票索引
    skip_idx = stock_df[stock_df['股票代码'].isin(skipped_list)].index
    # 按索引删除行
    stock_df = stock_df.drop(skip_idx, axis=0)
    print(f"\n🔍 统计结果: 共 {len(stock_codes)} 只股票 | {len(skipped_list)} 只已跳过 | {len(to_process)} 只需要处理\n")

    max_workers = min(20, len(stock_df))  # 设置最大线程数为10或股票数量，取较小值
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = []
        i = 0
        for index, row in stock_df.iterrows():
            symbol = row["股票代码"]  # 获取当前行的股票代码
            stock_name = row["股票名称"]  # 获取当前行的股票名称
            # 提交任务到线程池
            future = executor.submit(
                process_single_stock_without_check,
                symbol=symbol,
                stock_name=stock_name,
                start_date=start_date,
                end_date=end_date
            )
            futures.append((future, symbol))
            i = i + 1
            print(f"已提交{symbol}的处理任务,当前进度为{i}/{len(to_process)} \n")

        print("所有任务已提交，开始处理...")

        # 处理完成的任务
        success_list = []
        fail_list = []
        for future, symbol in futures:
            try:
                symbol_result, status, message = future.result()
                if status == "success":
                    success_list.append(symbol_result)
                    print(f"✅ {symbol_result} 处理成功")
                else:
                    fail_list.append(symbol_result)
                    print(f"❌ {symbol_result} 处理失败: {message}")
            except Exception as e:
                fail_list.append(symbol)
                print(f"❌ {symbol} 任务处理异常: {str(e)}")

        # 添加睡眠延迟
        sleepTime = random.uniform(0.2, 0.4)
        print(f"休息{sleepTime:.2f}秒后继续\n")
        time.sleep(sleepTime)
        end = time.perf_counter()
        print(f"目前总耗时: {end - start:.2f}秒")
        print("\n" + "=" * 50)

    # 处理完成的任务
    for future in as_completed(futures):
        try:
            symbol, status, message = future.result()
            if status == "success":
                success_list.append(symbol)
            else:
                fail_list.append(symbol)
        except Exception as e:
            print(f"❌ 任务处理异常: {str(e)}")
            fail_list.append(symbol)

    # 打印总结
    print("\n" + "=" * 50)
    print("📊 多线程任务完成！")
    print(f"✅ 成功处理: {len(success_list)} 只股票")
    print(f"⏩ 跳过处理: {len(skipped_list)} 只股票（已存在）")
    print(f"❌ 失败处理: {len(fail_list)} 只股票")

    if skipped_list:
        print("\n跳过的股票代码: " + ", ".join(skipped_list))
    if fail_list:
        print("\n失败的股票代码: " + ", ".join(fail_list))

    # 保存处理摘要
    summary_data = {
        '股票代码': success_list + skipped_list + fail_list,
        '状态': ['成功'] * len(success_list) + ['已跳过'] * len(skipped_list) + ['失败'] * len(fail_list),
        '备注': ['新获取'] * len(success_list) + ['文件已存在'] * len(skipped_list) + ['获取失败'] * len(fail_list)
    }

    summary = pd.DataFrame(summary_data)
    summary_path = os.path.join("../gp_xlsx", "处理摘要.xlsx")
    summary.to_excel(summary_path, index=False)
    print(f"\n📝 处理摘要已保存: {os.path.abspath(summary_path)}")
    print("=" * 50)


if __name__ == "__main__":
    # 自动安装依赖
    try:
        import fake_useragent
    except ImportError:
        print("⚠️ 正在安装 fake-useragent 库...")
        import subprocess

        subprocess.run(['pip', 'install', 'fake_useragent'], check=True)
        print("✅ 依赖安装完成，重新运行程序")
        exit(0)

    main()